#! python3
# -*- coding: UTF-8 -*-
# @Time    : 2025/05/26
# @Author  : AndrewWay
# @version : 1.0
# @notes   ：

"""
数据库连接
"""

import pymysql
from pymysql import Connection
from pymysql.cursors import Cursor
from setting import DatabaseConfig
from sqlalchemy import create_engine, QueuePool
from sqlalchemy.orm import sessionmaker, Session
from loguru import logger


engine = create_engine(DatabaseConfig.SQLALCHEMY_URI,
                       pool_pre_ping=True,
                       # pool=QueuePool,
                       pool_size=DatabaseConfig.pool_size,
                       max_overflow=DatabaseConfig.max_overflow,
                       pool_recycle=DatabaseConfig.pool_recycle,
                       pool_timeout=DatabaseConfig.pool_timeout,
                       echo=True,  # 开发时开启，打印sql语句
                       )

session_maker = sessionmaker(bind=engine, expire_on_commit=False, autocommit=False, autoflush=False,)


def get_session() -> Session:
    """
    获取数据库连接
    :return:
    """
    session = session_maker()
    try:
        yield session
    finally:
        try:
            session.close()
        except:
            pass


def shutdown_engine():
    """
    关闭数据库连接池
    :return:
    """
    try:
        engine.dispose()
    except Exception as e:
        error_msg = "关闭数据库连接池失败，请检查！"
        logger.warning(error_msg)


class DB:
    """
    数据库链接
    """

    def _get_connect(self):
        self.conn = pymysql.connect(host=DatabaseConfig.host,
                                    port=DatabaseConfig.port,
                                    user=DatabaseConfig.username,
                                    password=DatabaseConfig.password,
                                    database=DatabaseConfig.database,
                                    charset=DatabaseConfig.charset,
                                    )
        self.csr = self.conn.cursor()

    def __init__(self):
        self.conn: Connection = None
        self.csr: Cursor = None
        self._get_connect()

    def query_one(self, sql, values=None):
        self.csr.execute(sql, args=values)
        return self.csr.fetchone()

    def query(self, sql, values=None):
        self.csr.execute(sql, args=values)
        return self.csr.fetchall()

    def exec_many(self, sql, values):
        results = self.csr.executemany(sql, tuple(values))
        self.conn.commit()
        return results

    def exec_one(self, sql, values=None):
        results = self.csr.execute(sql, values)
        self.conn.commit()
        return results

    def close_db(self):
        try:
            self.csr.close()
            self.conn.close()
        except Exception as e:
            logger.warning('close connection error ,result:{}'.format(e))

